from typing import Iterable, Set, Collection, Tuple
from collections import defaultdict

from chernc.treediff.gumtree.tree.tree import TreeLike

class HashBasedMapper:
    def __init__(self) -> None:
        # 使用 defaultdict 初始化映射，默认值为一个 Pair 对象
        self._mappings: dict[int, Tuple[Set[TreeLike], Set[TreeLike]]] = defaultdict(lambda: (set(), set()))

    def add_srcs(self, srcs: Collection[TreeLike]) -> None:
        for t in srcs:
            self.add_src(t)

    def add_dsts(self, dsts: Collection[TreeLike]) -> None:
        for t in dsts:
            self.add_dst(t)

    def add_src(self, src: TreeLike) -> None:
        # 根据 src 的 hash 值将其加入到映射中
        self._mappings[src.get_metrics().hash][0].add(src)

    def add_dst(self, dst: TreeLike) -> None:
        # 根据 dst 的 hash 值将其加入到映射中
        self._mappings[dst.get_metrics().hash][1].add(dst)

    def unique(self) -> Iterable[Tuple[Set[TreeLike], Set[TreeLike]]]:
        return (value for value in self._mappings.values()
                if len(value[0]) == 1 and len(value[1]) == 1)

    def ambiguous(self) -> Iterable[Tuple[Set[TreeLike], Set[TreeLike]]]:
        return (value for value in self._mappings.values()
                if (len(value[0]) > 1 and len(value[1]) >= 1)
                or (len(value[0]) >= 1 and len(value[1]) > 1))

    def unmapped(self) -> Iterable[Tuple[Set[TreeLike], Set[TreeLike]]]:
        return (value for value in self._mappings.values()
                if len(value[0]) == 0 or len(value[1]) == 0)

    def is_src_mapped(self, src: TreeLike) -> bool:
        return len(self._mappings[src.get_metrics().hash][1]) > 0

    def is_dst_mapped(self, dst: TreeLike) -> bool:
        return len(self._mappings[dst.get_metrics().hash][0]) > 0

